#!/usr/bin/env Rscript
if (!require("librarian", quietly = TRUE))
  install.packages("librarian")
library(librarian)
shelf(argparse, readr, tibble, dplyr, stringr, tidyr, purrr, crayon)

ena.api <- 'https://www.ebi.ac.uk/ena/portal/api/filereport?accession=PRJNA958247&result=read_run&fields=study_accession,run_accession,fastq_bytes,fastq_md5,fastq_ftp,fastq_aspera,sample_alias,sample_title&format=tsv&download=true&limit=0'

# to use in shell
parser <- ArgumentParser()

parser$add_argument("-i", "--input", nargs = "+", help = "Input ENA tsv file or ENA project accession, REQUIRED")
parser$add_argument("-k", "--keep_ends", default = "0", type = "integer",
  help = "Which end of read to keep? 0 to keep both ends [0,1,2]"
)
parser$add_argument("-n", "--new_csv", default = TRUE, action = "store_true",
                    help = "Output homemade-style csv instead of nfcore-style [default]"
)
parser$add_argument("-o", "--old_csv", dest = "new_csv", action = "store_false",
                    help = "Output nfcore-style csv"
)
parser$add_argument("-e", "--tenx5p_mode", default = FALSE, action = "store_true",
                    help = "Exchange R1 with R2 for 10x 5' protocol? (require -n)"
)

input_args <- parser$parse_args()

# get input file path ------
input_ena <- input_args$input

## auto convert input to api url if accession not file is detected 
if (str_ends(input_ena, '\\d')) {
  message("INFO: accession detected. Download meta data from ENA api...")
  ena.acc <- input_ena
  input_ena <- ena.api |>
    str_replace('PRJNA\\d+',input_ena)
  
  download.file(input_ena, str_c(ena.acc, '.tsv'))
}

ena <- read_tsv(input_ena, show_col_types = FALSE) |>
  arrange(sample_title)

# determine if the data is dual-end or single-end --------
is_dual_end <- ena$fastq_ftp[1L] |>
  str_detect(";")

# nothing will change in single-end data
long_ena <- ena |>
  separate_longer_delim(contains('fastq'), ';') |>
  filter(!is.na(fastq_ftp))

long_ena <- long_ena |>
  mutate(end = str_extract(fastq_ftp, '.(?=\\.fastq\\.gz)'),
         fastq_bytes = as.numeric(fastq_bytes))

# make wide_ena tibble for nfcore input --------
if (is_dual_end) {
  long_ena <- long_ena |>
    mutate(end = str_extract(fastq_ftp, '.(?=\\.fastq\\.gz)'))
  
  wide_ena <- ena |>
    separate(fastq_ftp, into = c("fastq_1", "fastq_2"), sep = ";")
} else {
  # auto set strand mode to forward if is SE data
  input_args$keep_ends <- "1"
  long_ena <- long_ena |> mutate(end = '1')
  message("INFO: single-end data detected. Override with forward-end mode.")
  wide_ena <- ena |>
    mutate(fastq_1 = fastq_ftp, fastq_2 = '')
}

# pick different ends data in need ------
if (input_args$keep_ends == "0") {
  message("INFO: keep dual-end data.")
  wide_ena <- wide_ena |>
    mutate(strandedness = "unstranded")
} else if (input_args$keep_ends == "2") {
  message("INFO: keep ONLY reverse-end data.")
  long_ena <- long_ena |> filter(end == '2')
  wide_ena <- wide_ena |>
    mutate(strandedness = "reverse", fastq_1 = fastq_2, fastq_2 = '')
} else {
  message("INFO: keep ONLY forward-end data.")
  long_ena <- long_ena |> filter(end != '2')
  wide_ena <- wide_ena |>
    mutate(strandedness = "forward", fastq_2 = '')
}

# report some statistics summary ---------
sum_size <- round(sum(long_ena$fastq_bytes) / 1024^3, 2)
ftp_list <- long_ena$fastq_ftp
ascp_list <- long_ena$fastq_aspera

fct <- length(ftp_list)

message(str_c("Total size in GB: ", sum_size))

message(str_c("File counts: ", fct))

message(str_c("Average size in GB: ", round(sum_size / fct, 2)))

if (sum_size * 5 > 500) {
  inflate_warn <- str_c("Result data size may inflate up to ~ ", round(sum_size * 5, 2), ' GB! Please ensure enough free space or consider split the pipeline runs.')
  warning(blue(inflate_warn))
}

# write ftp & aspera links, md5 and nfcore input sheet --------
tibble(ftp_link = ftp_list) |>
  mutate(ftp_link = str_glue("ftp://{ftp_link}")) |>
  write_tsv("ftp_links.txt", col_names = FALSE)

tibble(ascp_link = ascp_list) |>
  mutate(ascp_link = str_glue("era-fasp@{ascp_link}")) |>
  write_tsv("ascp_links.txt", col_names = FALSE)

# raise warning if no md5 provided
if (some(colnames(long_ena), \(x)str_detect(x,'md5'))) {
  message('md5 detected. Good practice.')
  long_ena |>
    mutate(md5 = fastq_md5, fastq = str_remove(fastq_ftp, ".+/"), .keep = 'none') |>
    write_tsv("md5.txt", col_names = FALSE)
} else {
  warning(bgWhite$red('No md5! Beware of file corruption!'))
}

wdir <- getwd()

# replace everything before the last '/' with work dir path
# xxx:/xxx/xxx/SRRxxx.fq.gz -> /work/dir/SRRxxx.fq.gz
nf_csv <- wide_ena |>
  mutate(
    fastq_2 = str_replace(fastq_2, ".+(?=/)", wdir),
    fastq_1 = str_replace(fastq_1, ".+(?=/)", wdir),
    sample = make.names(sample_title) |> str_replace_all('\\.','_'),
    strandedness = strandedness,
    .keep = "none"
  ) |>
  relocate(sample) 

nf_csv |>
  write_csv("nfcore_input.csv")

if (input_args$tenx5p_mode & !input_args$new_csv) {
  warning(blue("new_csv not enabled, 10x 5' mode will not work."))
}

if (input_args$new_csv) {
  if (input_args$tenx5p_mode) {
    message('Swap ends for 10x 5 prime protocol...')
    nf_csv <- nf_csv |>
      mutate(fastq_1.1 = fastq_1, fastq_1 = fastq_2, fastq_2 = fastq_1.1,
             fastq_1.1 = NULL)
  }
  nf_csv |>
    mutate(readgroup = sample, strandedness = NULL) |>
    relocate(sample, readgroup, fastq_1) |>
    write_csv("nxf_input.csv")
}